import av
from PIL import Image

from app.utils.file_util import watermark_video_path


# 返回没有音频
def add_watermark(input_video, watermark_image, position):
    output_video = watermark_video_path(input_video)
    with av.open(input_video) as in_container, av.open(output_video, mode='w') as out_container:
        watermark = Image.open(watermark_image)
        # 添加输出流
        in_video_stream = in_container.streams.video[0]
        video_stream = out_container.add_stream('h264', rate=in_container.streams.video[0].average_rate)
        video_stream.width = in_video_stream.width
        video_stream.height = in_video_stream.height
        video_stream.pix_fmt = in_video_stream.pix_fmt

        for frame in in_container.decode(video=0):
            img = frame.to_image()
            img.paste(watermark, position, watermark)
            new_frame = av.VideoFrame.from_image(img)
            for packet in video_stream.encode(new_frame):
                out_container.mux(packet)

        for packet in video_stream.encode():
            out_container.mux(packet)
    return output_video


if __name__ == '__main__':
    add_watermark('input.mp4', 'output.mp4', 'watermark.png')
